package org.finesys.chat.core.base.client;

import com.fasterxml.jackson.databind.JsonNode;
import com.zhipu.oapi.ClientV4;
import com.zhipu.oapi.service.v4.embedding.EmbeddingApiResponse;
import com.zhipu.oapi.service.v4.embedding.EmbeddingRequest;
import com.zhipu.oapi.service.v4.image.CreateImageRequest;
import com.zhipu.oapi.service.v4.image.Image;
import com.zhipu.oapi.service.v4.image.ImageApiResponse;
import com.zhipu.oapi.service.v4.model.*;
import lombok.extern.slf4j.Slf4j;
import org.finesys.ai.core.chat.ChatCompletionResponse;
import org.finesys.ai.core.completion.CompletionChoice;
import org.finesys.ai.core.completion.CompletionResponse;
import org.finesys.ai.core.embedding.Embedding;
import org.finesys.ai.core.embedding.EmbeddingResponse;
import org.finesys.ai.core.image.GenerateImagesResponse;
import org.finesys.ai.core.image.ImageData;
import org.finesys.ai.core.shared.Usage;
import org.finesys.chat.core.base.chat.ChatRequest;
import org.finesys.chat.core.base.chat.ChatToolCall;
import org.finesys.chat.core.base.factory.ChatModel;
import reactor.core.publisher.Flux;

import java.net.URI;
import java.util.ArrayList;
import java.util.Collections;
import java.util.List;
import java.util.UUID;
import java.util.concurrent.TimeUnit;

@Slf4j
public class ZhipuChatClient extends ChatClientAbstract {


    private ClientV4 clientV4;

    public ZhipuChatClient(ChatModel chatModel) {
        super(chatModel);
        clientV4 = new ClientV4.Builder(chatModel.getBaseUrl(), chatModel.getApiSecret())
                .setTokenKey(chatModel.getApiKey(), chatModel.getApiSecret())
                .networkConfig(chatModel.getCallTimeout(), chatModel.getConnectTimeout(), chatModel.getReadTimeout(), chatModel.getWriteTimeout(), TimeUnit.SECONDS).build();
    }


    @Override
    public void register(ChatModel chatModel) {
        super.register(chatModel);
        if (clientV4 != null) {
            throw new RuntimeException("clientV4 is not null");
        }
        clientV4 = new ClientV4.Builder(chatModel.getBaseUrl(), chatModel.getApiSecret())
                .setTokenKey(chatModel.getApiKey(), chatModel.getApiSecret())
                .networkConfig(chatModel.getCallTimeout(), chatModel.getConnectTimeout(), chatModel.getReadTimeout(), chatModel.getWriteTimeout(), TimeUnit.SECONDS).build();

    }

    @Override
    public void updateKey(ChatModel chatModel) {
        super.updateKey(chatModel);
        if (clientV4 == null) {
            throw new RuntimeException("clientV4 is null,can not update key");
        }
        clientV4 = null;
        clientV4 = new ClientV4.Builder(chatModel.getBaseUrl(), chatModel.getApiSecret())
                .setTokenKey(chatModel.getApiKey(), chatModel.getApiSecret())
                .networkConfig(chatModel.getCallTimeout(), chatModel.getConnectTimeout(), chatModel.getReadTimeout(), chatModel.getWriteTimeout(), TimeUnit.SECONDS).build();

    }

    @Override
    public void unregister() {
        super.unregister();
        clientV4 = null;
    }

    @Override
    public Flux<CompletionResponse> completion(ChatRequest chatRequest) {
        ModelApiResponse modelApiResponse = clientV4.invokeModelApi(buildeChatCompletionRequest(chatRequest));
        if (chatRequest.getStream()) {
            return Flux.from(modelApiResponse.getFlowable()).map(jsonNode -> {
                CompletionResponse.Builder builder = CompletionResponse.builder();
                builder.id(jsonNode.getId());
                builder.model(jsonNode.getModel());
                builder.created(jsonNode.getCreated().intValue());
                builder.usage(builderUsage(jsonNode.getUsage()));
                builder.choices(builderChoices(jsonNode.getChoices()));
                builder.systemFingerprint(jsonNode.getText());
                return builder.build();
            });
        }
        return Flux.just(builderCompletionResponse(modelApiResponse));
    }

    private CompletionResponse builderCompletionResponse(ModelApiResponse modelApiResponse) {
        CompletionResponse.Builder builder = CompletionResponse.builder();
        builder.id(modelApiResponse.getData().getId());
        builder.model(modelApiResponse.getData().getModel());
        builder.created(modelApiResponse.getData().getCreated().intValue());
        builder.usage(builderUsage(modelApiResponse.getData().getUsage()));
        builder.choices(builderChoices(modelApiResponse.getData().getChoices()));
        return builder.build();
    }

    private List<CompletionChoice> builderChoices(List<Choice> choices) {
        List<CompletionChoice> completionChoices = new ArrayList<>();
        choices.forEach(choice -> {
            CompletionChoice.Builder completionChoiceBuilder = CompletionChoice.builder();
            completionChoiceBuilder.index(choice.getIndex().intValue());
            completionChoiceBuilder.finishReason(choice.getFinishReason());
            completionChoiceBuilder.text(choice.asText());
            completionChoices.add(completionChoiceBuilder.build());
        });
        return completionChoices;
    }

    private Usage builderUsage(com.zhipu.oapi.service.v4.model.Usage usage) {
        Usage.Builder usageBuilder = Usage.builder();
        usageBuilder.promptTokens(usage.getPromptTokens());
        usageBuilder.totalTokens(usage.getTotalTokens());
        usageBuilder.completionTokens(usage.getCompletionTokens());
        return usageBuilder.build();
    }

    private ChatCompletionRequest buildeChatCompletionRequest(ChatRequest chatRequest) {
        ChatCompletionRequest.ChatCompletionRequestBuilder<?, ?> builder = ChatCompletionRequest.builder();
        builder.model(chatRequest.getModel());
        builder.maxTokens(chatRequest.getMaxTokens());
        builder.temperature(chatRequest.getTemperature().floatValue());
        builder.topP(chatRequest.getTopP().floatValue());
        builder.stream(chatRequest.getStream());
        builder.messages(builderMessages(chatRequest.getMessages()));
        builder.stop(stopWords(chatRequest.getStopWords()));
        return builder.build();
    }

    private List<ChatMessage> builderMessages(List<org.finesys.chat.core.base.chat.ChatMessage> messages) {
        List<ChatMessage> chatMessages = new ArrayList<>();
        messages.forEach(message -> {
            ChatMessage builder = new ChatMessage();
            builder.setRole(message.getRole().getValue());
            builder.setContent(message.getContent());
            builder.setName(message.getName());
            builder.setTool_calls(builderZhipuToolCalls(message.getToolCalls()));
            chatMessages.add(builder);
        });
        return chatMessages;
    }

    private List<ToolCalls> builderZhipuToolCalls(List<ChatToolCall> toolCalls) {
        List<ToolCalls> toolCallsList = new ArrayList<>();
        toolCalls.forEach(toolCall -> {
            ToolCalls builder = new ToolCalls();
            builder.setId(toolCall.getId());
            builder.setType(toolCall.getType().name());
            builder.setFunction(builderZhipuFunctionCall(toolCall.getFunction()));
            toolCallsList.add(builder);
        });
        return toolCallsList;
    }

    private ChatFunctionCall builderZhipuFunctionCall(org.finesys.chat.core.base.chat.ChatFunctionCall function) {
        ChatFunctionCall builder = new ChatFunctionCall();
        builder.setName(function.getName());
        JsonNode jsonNode = org.finesys.commom.jackson.util.JacksonUtil.parseTree(function.getArguments());
        builder.setArguments(jsonNode);
        return builder;
    }

    @Override
    public Flux<ChatCompletionResponse> chatCompletion(ChatRequest chatRequest) {

        return super.chatCompletion(chatRequest);
    }

    private CreateImageRequest builderImageRequest(ChatRequest chatRequest) {
        CreateImageRequest.CreateImageRequestBuilder<?, ?> builder = CreateImageRequest.builder();
        builder.model(chatRequest.getModel());
        builder.size(chatRequest.getImageSize());
        builder.prompt(chatRequest.getUserMessage());
        builder.requestId(UUID.randomUUID().toString());
        return builder.build();
    }

    private GenerateImagesResponse builderGenerateImagesResponse(ChatRequest chatRequest, ImageApiResponse imageApiResponse) {
        GenerateImagesResponse.Builder builder = GenerateImagesResponse.builder();
        if (imageApiResponse.isSuccess()) {
            builder.data(builderImageData(imageApiResponse.getData().getData()));
        } else {
            builder.data(Collections.emptyList());
            log.error("image generation error: {}", imageApiResponse.getError().getMessage());
            throw new RuntimeException(imageApiResponse.getError().getMessage());
        }
        return builder.build();
    }

    private List<ImageData> builderImageData(List<Image> data) {
        List<ImageData> imageData = new ArrayList<>();
        data.forEach(image -> {
            ImageData.Builder imageDataBuilder = ImageData.builder();
            imageDataBuilder.url(URI.create(image.getUrl()));
            imageDataBuilder.b64Json(image.getB64Json());
            imageDataBuilder.revisedPrompt(image.getRevisedPrompt());
            imageData.add(imageDataBuilder.build());
        });
        return imageData;
    }

    @Override
    public Flux<GenerateImagesResponse> imageGeneration(ChatRequest chatRequest) {
        CreateImageRequest createImageRequest = builderImageRequest(chatRequest);
        ImageApiResponse imageApiResponse = clientV4.createImage(createImageRequest);
        GenerateImagesResponse generateImagesResponse = builderGenerateImagesResponse(chatRequest, imageApiResponse);
        return Flux.just(generateImagesResponse);
    }

    @Override
    public Flux<EmbeddingResponse> embedding(ChatRequest chatRequest) {
        EmbeddingApiResponse embeddingApiResponse = clientV4.invokeEmbeddingsApi(builderEmbeddingRequest(chatRequest));
        EmbeddingResponse embeddingResponse = builderEmbeddingResponse(embeddingApiResponse);
        return Flux.just(embeddingResponse);
    }

    private EmbeddingResponse builderEmbeddingResponse(EmbeddingApiResponse embeddingApiResponse) {
        EmbeddingResponse.Builder builder = EmbeddingResponse.builder();
        builder.model(embeddingApiResponse.getData().getModel());
        builder.usage(builderUsage(embeddingApiResponse.getData().getUsage()));
        builder.data(builderEmbindData(embeddingApiResponse.getData().getData()));

        return builder.build();
    }

    private List<Embedding> builderEmbindData(List<com.zhipu.oapi.service.v4.embedding.Embedding> data) {
        List<Embedding> embeddings = new ArrayList<>();
        data.forEach(embedding -> {
            Embedding.Builder embeddingBuilder = Embedding.builder();
            embeddingBuilder.index(embedding.getIndex());
            embeddingBuilder.embedding(embedding.getEmbedding());
            embeddings.add(embeddingBuilder.build());
        });
        return embeddings;
    }

    private EmbeddingRequest builderEmbeddingRequest(ChatRequest chatRequest) {
        EmbeddingRequest.EmbeddingRequestBuilder<?, ?> builder = EmbeddingRequest.builder();
        builder.model(chatRequest.getModel());
        builder.input(chatRequest.getInput());
        builder.requestId(UUID.randomUUID().toString());
        builder.dimensions(chatRequest.getDimensions());
        builder.userId(chatRequest.getUser());
        return builder.build();
    }
}
